import rospy
from arm_control.msg import JointControl
from arm_control.msg import JointInformation
import numpy as np
import time

class DualRobot:
    def __init__(self):
        print("init Dual Robot")
        rospy.init_node('dual', anonymous=True)
        rospy.Subscriber("/joint_information", JointInformation, self.callback_ji_left)
        rospy.Subscriber("/joint_information2", JointInformation, self.callback_ji_right)
        self.pub_left = rospy.Publisher('/pv_joint_control', JointControl, queue_size=10)
        self.pub_right = rospy.Publisher('/pv_joint_control2', JointControl, queue_size=10)
        self.left_pos = []
        self.right_pos = []

    def callback_ji_left(self,data):
        self.left_pos = data.joint_pos
    def callback_ji_right(self,data):
        self.right_pos = data.joint_pos

    def reset(self):
        step_time_ms = 50
        step_num = 100
        pos_left = [np.linspace(start=0, stop=0, num=step_num)]*7
        pos_right = [np.linspace(start=0, stop=0, num=step_num)]*7
        for i in range(7):
            pos_left[i] = np.linspace(start=self.left_pos[i], stop=0, num=step_num)
            pos_right[i] = np.linspace(start=self.right_pos[i], stop=0, num=step_num)

        data_left = JointControl()
        data_right = JointControl()
        for i in range(step_num):
            for j in range(7):
                data_left.joint_pos[j] = pos_left[j][i]
                data_right.joint_pos[j] = pos_right[j][i]
            self.pub_left.publish(data_left)
            self.pub_right.publish(data_right)
            time.sleep(step_time_ms/1000)



if __name__ == '__main__':
    robot = DualRobot()
    time.sleep(2)
    robot.reset()
